{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.0480, 0.0907, 0.0163,  ..., 0.0062, 0.0052, 0.0088],\n",
      "         [0.0178, 0.0221, 0.0092,  ..., 0.0206, 0.0158, 0.0080],\n",
      "         [0.0727, 0.0152, 0.0049,  ..., 0.0042, 0.0229, 0.0036],\n",
      "         ...,\n",
      "         [0.1803, 0.0335, 0.0159,  ..., 0.0563, 0.0206, 0.0083],\n",
      "         [0.0325, 0.0082, 0.0295,  ..., 0.0057, 0.0265, 0.0075],\n",
      "         [0.0191, 0.0232, 0.0773,  ..., 0.0121, 0.0460, 0.0269]]],\n",
      "       device='cuda:0', requires_grad=True)\n",
      "tensor([[[[0.0480],\n",
      "          [0.0907],\n",
      "          [0.0163],\n",
      "          ...,\n",
      "          [0.0062],\n",
      "          [0.0052],\n",
      "          [0.0088]],\n",
      "\n",
      "         [[0.0178],\n",
      "          [0.0221],\n",
      "          [0.0092],\n",
      "          ...,\n",
      "          [0.0206],\n",
      "          [0.0158],\n",
      "          [0.0080]],\n",
      "\n",
      "         [[0.0727],\n",
      "          [0.0152],\n",
      "          [0.0049],\n",
      "          ...,\n",
      "          [0.0042],\n",
      "          [0.0229],\n",
      "          [0.0036]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[0.1803],\n",
      "          [0.0335],\n",
      "          [0.0159],\n",
      "          ...,\n",
      "          [0.0563],\n",
      "          [0.0206],\n",
      "          [0.0083]],\n",
      "\n",
      "         [[0.0325],\n",
      "          [0.0082],\n",
      "          [0.0295],\n",
      "          ...,\n",
      "          [0.0057],\n",
      "          [0.0265],\n",
      "          [0.0075]],\n",
      "\n",
      "         [[0.0191],\n",
      "          [0.0232],\n",
      "          [0.0773],\n",
      "          ...,\n",
      "          [0.0121],\n",
      "          [0.0460],\n",
      "          [0.0269]]]], device='cuda:0', grad_fn=<UnsqueezeBackward0>)\n",
      "X.SIZE torch.Size([1, 32, 64, 32, 64])\n",
      "A.SIZE torch.Size([1, 32, 64, 32])\n",
      "B.size torch.Size([1, 32, 64, 1, 64])\n",
      "C.size torch.Size([1, 32, 64, 1, 64])\n",
      "A.SIZE torch.Size([1, 32, 32, 64])\n",
      "A_cum.SIZE torch.Size([1, 32, 32, 64])\n"
     ]
    }
   ],
   "source": [
    "from einops import rearrange, repeat\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "block_len=64\n",
    "torch.manual_seed(42)\n",
    "\n",
    "## Dimensions\n",
    "# Denoted (B, T, Q, D, P) in the paper\n",
    "batch, seqlen, chunk_size, dim, headdim = 1, 2048, 64, 2048, 64\n",
    "nheads = dim // headdim  # (H) in the paper\n",
    "ngroups = 1 # (G) in the paper\n",
    "dstate = 64  # (N) in the paper\n",
    "dtype = torch.float32\n",
    "device = \"cuda\"\n",
    "\n",
    "x = torch.randn(batch, seqlen, nheads, headdim, dtype=dtype, device=device)\n",
    "#print(x)\n",
    "#print(\"x.size\",x.size())  #x.size torch.Size([1, 2048, 32, 64])\n",
    "dt = F.softplus(torch.randn(batch, seqlen, nheads, dtype=torch.float32, device=device) - 4).requires_grad_()\n",
    "#print(dt)\n",
    "#print(\"dt.size\",dt.size())  #dt.size torch.Size([1, 2048, 32])\n",
    "A = (-torch.exp(torch.rand(nheads, dtype=torch.float32, device=device))).requires_grad_()\n",
    "#print(A)\n",
    "#print('A.SIZE',A.size())  #A.SIZE torch.Size([32])\n",
    "B = torch.randn(batch, seqlen, ngroups, dstate, dtype=dtype, device=device)  #B.size torch.Size([1, 2048, 1, 64])\n",
    "#print(\"B.size\",B.size())  \n",
    "C = torch.randn(batch, seqlen, ngroups, dstate, dtype=dtype, device=device)  # C.size torch.Size([1, 2048, 1, 64])\n",
    "#print(\"C.size\",C.size())\n",
    "D = torch.randn(nheads, dtype=dtype, device=device)  # D.size torch.Size([32])\n",
    "#print(\"D.size\",D.size())\n",
    "\n",
    "#print(x)\n",
    "print(dt)\n",
    "print(dt.unsqueeze(-1) )\n",
    "X=x*dt.unsqueeze(-1)   #X.SIZE torch.Size([1, 2048, 32, 64])\n",
    "#print(\"X.SIZE\",X.size()) \n",
    "#print(X)\n",
    "A=A*dt\n",
    "#print(\"A.SIZE\",A.size())  # A.SIZE torch.Size([1, 2048, 32])\n",
    "\n",
    "X, A, B, C = [rearrange(x, \"b (c l) ... -> b c l ...\", l=block_len) for x in (X, A, B, C)] \n",
    "print(\"X.SIZE\",X.size()) \n",
    "print(\"A.SIZE\",A.size())\n",
    "print(\"B.size\",B.size()) \n",
    "print(\"C.size\",C.size())\n",
    "\n",
    "'''X.SIZE torch.Size([1, 32, 64, 32, 64])\n",
    "A.SIZE torch.Size([1, 32, 64, 32])\n",
    "B.size torch.Size([1, 32, 64, 1, 64])\n",
    "C.size torch.Size([1, 32, 64, 1, 64])\n",
    "'''\n",
    "A = rearrange(A, \"b c l h -> b h c l\")\n",
    "A_cumsum = torch.cumsum(A, dim=-1)\n",
    "print(\"A.SIZE\",A.size())\n",
    "print(\"A_cum.SIZE\",A_cumsum.size())\n",
    "'''A.SIZE torch.Size([1, 32, 32, 64])\n",
    "A_cum.SIZE torch.Size([1, 32, 32, 64])'''\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
